from ryu.controller import ofp_event
from ryu.controller.handler import MAIN_DISPATCHER, DEAD_DISPATCHER
from ryu.controller.handler import set_ev_cls
from ryu.lib import hub
import simple_switch_stp_13

from datetime import datetime

import os
import numpy as np
from tensorflow.keras import models

import sys
sys.path.append('..')

from classifier.information_entropy import Entropy


os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

class CollectTrainingStatsApp(simple_switch_stp_13.SimpleSwitch13):
    def __init__(self, *args, **kwargs):
        super(CollectTrainingStatsApp, self).__init__(*args, **kwargs)
        self.datapaths = {}
        self.monitor_thread = hub.spawn(self.monitor)
        self.model = models.load_model('../saved_model/bpnn_model_input7.h5')    # 神经网络模型
        self.entropy = Entropy()    # 香农信息熵

    # Asynchronous message
    @set_ev_cls(ofp_event.EventOFPStateChange, [MAIN_DISPATCHER, DEAD_DISPATCHER])
    def state_change_handler(self, ev):
        datapath = ev.datapath
        if ev.state == MAIN_DISPATCHER:
            if datapath.id not in self.datapaths:
                self.logger.debug('register datapath: %016x', datapath.id)
                self.datapaths[datapath.id] = datapath

        elif ev.state == DEAD_DISPATCHER:
            if datapath.id in self.datapaths:
                self.logger.debug('unregister datapath: %016x', datapath.id)
                del self.datapaths[datapath.id]

    def monitor(self):
        while True:
            for dp in self.datapaths.values():
                self.request_stats(dp)
            hub.sleep(10)

    def request_stats(self, datapath):
        self.logger.debug('######## send stats request: %016x', datapath.id)

        parser = datapath.ofproto_parser

        req = parser.OFPFlowStatsRequest(datapath)
        datapath.send_msg(req)

    # 被动防御ddos攻击
    def punish(self, datapath, ip):
        ofproto = datapath.ofproto
        parser = datapath.ofproto_parser

        match = parser.OFPMatch(ipv4_src=(ip, '255.255.255.0'))
        actions = [parser.OFPActionOutput(4)]
        self.add_flow(datapath, 11, match, actions, idle=10, hard=10)

    @set_ev_cls(ofp_event.EventOFPFlowStatsReply, MAIN_DISPATCHER)
    def _flow_stats_reply_handler(self, ev):

        body = ev.msg.body
        for stat in sorted([flow for flow in body if (flow.priority == 1)],
                           key=lambda flow:(
                               flow.match['eth_type'],
                               flow.match['ipv4_src'],
                               flow.match['ipv4_dst'],
                               flow.match['ip_proto']
                           )):

            ip_src = stat.match['ipv4_src']
            ip_dst = stat.match['ipv4_dst']

            ip_proto = stat.match['ip_proto']    # feature 1
            flow_duration_sec = stat.duration_sec    # feature 2
            packet_count = stat.packet_count    # feature 3
            # feature 4
            try:
                packet_count_per_second = stat.packet_count / stat.duration_sec
            except:
                packet_count_per_second = 0

            byte_count = stat.byte_count    # feature 5
            # feature 6
            try:
                byte_count_per_second = stat.byte_count / stat.duration_sec
            except:
                byte_count_per_second = 0

            self.entropy.collect(ip_dst)
            entro_val = self.entropy.value  # feature 7

            # 判断是否是ddos攻击流量
            x = np.array([ip_proto,
                          flow_duration_sec,
                          packet_count,
                          packet_count_per_second,
                          byte_count,
                          byte_count_per_second,
                          entro_val]).reshape((1, -1))

            # 先进行信息熵判断，正常流量熵值范围为1.3~2.3
            if (self.entropy.value <= 1.3 or self.entropy.value >= 2.3):    # 不在范围内可能为异常流量，需要进一步判断
                print('### Abnormal traffic.  Information entropy: ', self.entropy.value)
                # 判断异常攻击
                pre = self.model.predict(x)
                if pre >= 0.21008658:    # 神经网络判断是否为攻击流量
                    print('\n-------- ddos attack --------')
                    print('  ip src: ', ip_src)
                    print('  ip dst: ', ip_dst)
                    print('-----------------------------')
                    # 不要惩罚服务器
                    # if ip_src != '10.0.0.5': self.punish(datapath=ev.msg.datapath, ip=ip_src)
                else:
                    print('\n-------- false alarm --------')
                    print('-----------------------------')
            else:
                print('### Healthy traffic.  Information entropy: ', self.entropy.value)
